from .UDP消息 import UDP消息, UDP消息类型
from .UDP协议 import UDP协议
from .UDP会话 import UDP会话
from .异常 import *
import asyncio


class UDP客户端:
    def __init__(self, 协议: UDP协议, 会话ID: int) -> None:
        self.协议 = 协议
        self.会话: UDP会话 = UDP会话(self, 协议, 会话ID)
        self.loop = asyncio.get_running_loop()
        self.消息收取任务 = self.loop .create_task(self.收消息任务())

    async def 收消息任务(self):
        while not self.协议.链接关闭:
            data = await self.协议.read()
            if not data:
                raise UDP链接已关闭()
            消息 = UDP消息.拆包(*data)
            try:
                self.loop.create_task(self.会话.收消息(消息))
            except UDP会话已关闭:
                self.协议.链接关闭 = True
                return
            except UDP链接已关闭:
                return

    def 关闭会话(self, 会话ID: int = 0):
        self.协议.链接关闭 = True
        self.消息收取任务.cancel()


async def open_connection(
        服务器地址: str, 服务器端口: int,
        超时时间: float = 5.0, 重连次数=3):
    loop = asyncio.get_running_loop()
    传输, 协议 = await loop.create_datagram_endpoint(
        lambda: UDP协议(),
        remote_addr=(服务器地址, 服务器端口)
    )

    请求消息 = UDP消息(UDP消息类型.请求建立链接, 0)
    for i in range(重连次数):
        if not 协议.write(请求消息()):
            raise UDP链接已关闭()
        try:
            data = await 协议.read(超时时间)
            if not data:
                raise UDP链接已关闭()
            消息 = UDP消息.拆包(*data)
            if 消息.消息类型 == UDP消息类型.链接成功:
                client = UDP客户端(协议, 消息.会话ID)
                if not 协议.write(消息.创建确认消息()()):
                    raise UDP链接已关闭()
                return client
        except 超时:
            print('超时')
